# load_dicom： 加载dicom文件，得到多个模态的序列【【name，itk-data】，【name-1，itk-data-1】】
# resample_standard_list 重采样itk list到相同尺寸，配准。
# resample_target_data 配准

# 此版本为分割过程中保存mask的版本

import torch.nn as nn 
from segmentation.segment import run_segmentation
from segmentation.model import load_segment_model
from segmentation.segment import get_resampled_target_data
from preprocessing.unzip import unzip_downloaded_dicom, get_all_modality_dir
from preprocessing.load_dicom import load_dicom, load_single_modality_dicom
import SimpleITK as sitk 
from postprocessing.utils import get_bboxes, save_bboxes
from .seg_result_to_dicom import save_segmentation_to_dicom
from Flask_run import db
from database.models.mri import Mri
from database.models.study_id import Study_id
import numpy as np

from pyorthanc import find, Orthanc
orthanc = Orthanc('http://localhost:8042',
                  username='orthanc', password='orthanc')

def save_segmentation_to_dicom_api(all_modality_dir, segmentation_data, study_id):
    used_modality_dir = None
    for modality_dir in all_modality_dir:
        # if "t1_fl2d_tra" in modality_dir:
        # 上面这行已修改为下面这行
        if "t1" in modality_dir and 'tra' in modality_dir:
          used_modality_dir = modality_dir
    save_segmentation_to_dicom(used_modality_dir, segmentation_data, study_id)
    
    
def check(dir_path):
  reader = sitk.ImageSeriesReader()
  # 利用读取器得到目录下所包含的序列ID，比如图中的CT WB 5.0 B31f、PET WB (AC)和ThorRoutine 2.0 B40f
  series_IDs = reader.GetGDCMSeriesIDs(dir_path)
  if len(series_IDs) != 1:
     return 0
  else:
     return 1
  
def save_t1_to_nii_api(all_modality_dir, PN):
    used_modality_dir = None
    for modality_dir in all_modality_dir:
        # if "t1_fl2d_tra" in modality_dir:
        # 上面这行已修改为下面这行
        if "t1" in modality_dir and 'tra' in modality_dir:
          used_modality_dir = modality_dir
    save_t1_nii(used_modality_dir,PN) 
       
def save_t1_nii(t1_dir,PN):

    reader = sitk.ImageSeriesReader()
    series_IDs = reader.GetGDCMSeriesIDs(t1_dir)
    series = series_IDs[0]
    dcm_files = reader.GetGDCMSeriesFileNames(t1_dir, series)
    reader.SetFileNames(dcm_files)
    image = reader.Execute()
    image_data = sitk.GetArrayFromImage(image)
    
    print("image_data.shape",image_data.shape)  
    sitk.WriteImage(image,f'./save_t1_nii/{PN}_t1.nii.gz')   

def getPN(study_id):
    print('-------------------------------------------')
    ort_study_id = db.session.query(Study_id.orthanc_id).filter(Study_id.study_id == study_id).first()
    ort_study_id = ort_study_id[0]
    print(ort_study_id)
    patients = find(
          orthanc=orthanc,
          study_filter=lambda s: s.id_ == ort_study_id
      )
    for patient in patients:
      for study in patient.studies:
        if study.id_ == ort_study_id:
            PN=study.patient_information['PatientName']
            print(PN)
            return PN
            
    



def segmentation_api(zip_from_orthanc_path, study_id):
    model = load_segment_model()
    # print(model)
    # 对下载后的数据进行分析，得到每个模态的目录
    all_files = unzip_downloaded_dicom(zip_from_orthanc_path)
    out = get_all_modality_dir(all_files)
    # print(out)
    ## 对每个模态进行加载，得到itk-list
    itk_list = []
    for d in out:
        status=check(d)
        if status:
          single_dicom = load_single_modality_dicom(d)
          itk_list.append(single_dicom)
    # 运行分割算法
    try:
        pred = run_segmentation(itk_list, model=model).cpu().numpy()
        print(f"segmentation shape is {pred.shape}")

        save_segmentation_to_dicom_api(out, pred, study_id)
        print(f"分割dicom保存成功")
    except (ValueError):
        db.session.query(Mri).filter_by(study_id=study_id).update({'available': False})
        db.session.commit()
        return 5000
    except (AssertionError):
        db.session.query(Mri).filter_by(study_id=study_id).update({'available': False})
        db.session.commit()
        return 10000
    # except:
    #     print("未知异常")
    #     db.session.query(Mri).filter_by(study_id=study_id).update({'available': False})
    #     db.session.commit()
    #     return 40000

    # mask保存为nii.gz
    PN=getPN(study_id)
    pred=pred.astype(np.float32)
    mask = sitk.GetImageFromArray(pred)
    sitk.WriteImage(mask,f'./save_seg_nii/{PN}_mask.nii.gz')
    # sitk.WriteImage(mask,f'./save_seg_nii/{study_id}_mask.nii.gz')
    # t1保存为nii.gz
    save_t1_to_nii_api(out, PN)


    resample_data_arr, _ = get_resampled_target_data(itk_list)
    bboxes = get_bboxes(three_modalities_arr=resample_data_arr, seg_out=pred)
    save_bboxes(bboxes, save_dir=f"../Dicom_download/bboxes_{study_id}")

    print(f"boxes saved successfully: bboxes_{study_id}")

    return 20000

